__author__ = 'Heiko Wundram <me@modelnine.org>'
__version__ = '0.2'
__revision__ = '3'
__date__ = '2006-01-20'
import intset
import socket

class IP4Range(intset.IntSet):
    _MINIP4 = 0
    _MAXIP4 = 4294967295L
    _UNITYTRANS = ''.join([ chr(n) for n in range(256) ])
    _IPREMOVE = '0123456789.'

    def __init__(self, *args):
        if ((len(args) == 1) and isinstance(args[0], IP4Range)):
            super(IP4Range, self).__init__(args[0])
            return 
        args = list(args)
        for i in range(len(args)):
            argval = args[i]
            if isinstance(argval, str):
                if ('<->' in argval):
                    args[i] = self._parseRange(*argval.split('<->', 1))
                    continue
                elif ('/' in argval):
                    args[i] = self._parseMask(*argval.split('/', 1))
                else:
                    args[i] = self._parseAddrRange(argval)
            elif isinstance(argval, tuple):
                if (len(tuple) != 2):
                    raise ValueError('Tuple is of invalid length.')
                (addr1, addr2,) = argval
                if isinstance(addr1, str):
                    addr1 = self._parseAddrRange(addr1)[0]
                elif not isinstance(addr1, (int, long)):
                    raise TypeError('Invalid argument.')
                if isinstance(addr2, str):
                    addr2 = self._parseAddrRange(addr2)[1]
                elif not isinstance(addr2, (int, long)):
                    raise TypeError('Invalid argument.')
                args[i] = (addr1, addr2)
            else:
                if isinstance(argval, (int, long)):
                    continue
                raise TypeError('Invalid argument.')
                continue

        super(IP4Range, self).__init__(min=self._MINIP4, max=self._MAXIP4, *args)



    def _parseRange(self, addr1, addr2):
        (naddr1, naddr1len,) = _parseAddr(addr1)
        (naddr2, naddr2len,) = _parseAddr(addr2)
        if (naddr2len < naddr1len):
            naddr2 += (naddr1 & (((1 << ((naddr1len - naddr2len) * 8)) - 1) << (naddr2len * 8)))
            naddr2len = naddr1len
        elif (naddr2len > naddr1len):
            raise ValueError('Range has more dots than address.')
        naddr1 <<= ((4 - naddr1len) * 8)
        naddr2 <<= ((4 - naddr2len) * 8)
        naddr2 += ((1 << ((4 - naddr2len) * 8)) - 1)
        return (naddr1, naddr2)



    def _parseMask(self, addr, mask):
        (naddr, naddrlen,) = _parseAddr(addr)
        naddr <<= ((4 - naddrlen) * 8)
        try:
            if not mask:
                masklen = 0
            else:
                masklen = int(mask)
            if not 0 <= masklen <= 32:
                raise ValueError
        except ValueError:
            try:
                mask = _parseAddr(mask, False)
            except ValueError:
                raise ValueError("Mask isn't parseable.")
            remaining = 0
            masklen = 0
            if not mask:
                masklen = 0
            else:
                while not (mask & 1):
                    remaining += 1

                while (mask & 1):
                    mask >>= 1
                    masklen += 1

                if ((remaining + masklen) != 32):
                    raise ValueError("Mask isn't a proper host mask.")
        naddr1 = (naddr & (((1 << masklen) - 1) << (32 - masklen)))
        naddr2 = ((naddr1 + (1 << (32 - masklen))) - 1)
        return (naddr1, naddr2)



    def _parseAddrRange(self, addr):
        (naddr, naddrlen,) = _parseAddr(addr)
        naddr1 = (naddr << ((4 - naddrlen) * 8))
        naddr2 = (((naddr << ((4 - naddrlen) * 8)) + (1 << ((4 - naddrlen) * 8))) - 1)
        return (naddr1, naddr2)



    def _int2ip(self, num):
        rv = []
        for i in range(4):
            rv.append(str((num & 255)))
            num >>= 8

        return '.'.join(reversed(rv))



    def iteraddresses(self):
        for v in super(IP4Range, self).__iter__():
            yield self._int2ip(v)




    def iterranges(self):
        for r in self._ranges:
            if ((r[1] - r[0]) == 1):
                yield self._int2ip(r[0])
            else:
                yield ('%s-%s' % (self._int2ip(r[0]), self._int2ip((r[1] - 1))))




    def itermasks(self):
        for r in self._ranges:
            for v in self._itermasks(r):
                yield v





    def _itermasks(self, r):
        ranges = [r]
        while ranges:
            cur = ranges.pop()
            curmask = 0
            while True:
                curmasklen = (1 << (32 - curmask))
                start = (((cur[0] + curmasklen) - 1) & (((1 << curmask) - 1) << (32 - curmask)))
                if ((start >= cur[0]) and ((start + curmasklen) <= cur[1])):
                    break
                else:
                    curmask += 1

            yield ('%s/%s' % (self._int2ip(start), curmask))
            if (cur[0] < start):
                ranges.append((cur[0], start))
            if (cur[1] > (start + curmasklen)):
                ranges.append(((start + curmasklen), cur[1]))



    __iter__ = iteraddresses

    def __repr__(self):
        rv = []
        for (start, stop,) in self._ranges:
            if ((stop - start) == 1):
                rv.append(('%r' % (self._int2ip(start))))
            else:
                rv.append(('(%r,%r)' % (self._int2ip(start), self._int2ip((stop - 1)))))

        return ('%s(%s)' % (self.__class__.__name__, ','.join(rv)))




def _parseAddr(addr, lookup = True):
    if (lookup and addr.translate(IP4Range._UNITYTRANS, IP4Range._IPREMOVE)):
        try:
            addr = socket.gethostbyname(addr)
        except socket.error:
            raise ValueError('Invalid Hostname as argument.')
    naddr = 0
    for (naddrpos, part,) in enumerate(addr.split('.')):
        if (naddrpos >= 4):
            raise ValueError('Address contains more than four parts.')
        try:
            if not part:
                part = 0
            else:
                part = int(part)
            if not 0 <= part < 256:
                raise ValueError
        except ValueError:
            raise ValueError('Address part out of range.')
        naddr <<= 8
        naddr += part

    return (naddr, (naddrpos + 1))



def ip2int(addr, lookup = True):
    return _parseAddr(addr, lookup=lookup)[0]


if (__name__ == '__main__'):
    x = IP4Range('172.22.162.250/24')
    y = IP4Range('172.22.162.250', '172.22.163.250', '172.22.163.253<->255')
    print x
    for val in x.itermasks():
        print val

    for val in y.itermasks():
        print val

    for val in (x | y).itermasks():
        print val

    for val in (x ^ y).iterranges():
        print val

    for val in x:
        print val


